{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Checkpoint and Resume Training for Semantic Segmentation Model\n",
    "\n",
    "[![image](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/opengeos/geoai/blob/main/docs/examples/load_model_checkpoint.ipynb)\n",
    "\n",
    "This notebook tests the new checkpoint loading and resume training functionality for the `train_segmentation_model` function. It demonstrates how to:\n",
    "\n",
    "1. Train a model for a few epochs\n",
    "2. Stop training and save a checkpoint\n",
    "3. Resume training from the checkpoint\n",
    "4. Load only model weights without resuming training state\n",
    "\n",
    "## Install packages\n",
    "To use the new functionality, ensure the required packages are installed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %pip install geoai-py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import geoai\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download sample data\n",
    "\n",
    "We'll use the same dataset as the main segmentation example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_raster_url = (\n",
    "    \"https://huggingface.co/datasets/giswqs/geospatial/resolve/main/naip_rgb_train.tif\"\n",
    ")\n",
    "train_vector_url = \"https://huggingface.co/datasets/giswqs/geospatial/resolve/main/naip_train_buildings.geojson\"\n",
    "test_raster_url = (\n",
    "    \"https://huggingface.co/datasets/giswqs/geospatial/resolve/main/naip_test.tif\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_raster_path = geoai.download_file(train_raster_url)\n",
    "train_vector_path = geoai.download_file(train_vector_url)\n",
    "test_raster_path = geoai.download_file(test_raster_url)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create training data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out_folder = \"checkpoint_test\"\n",
    "tiles = geoai.export_geotiff_tiles(\n",
    "    in_raster=train_raster_path,\n",
    "    out_folder=out_folder,\n",
    "    in_class_data=train_vector_path,\n",
    "    tile_size=512,\n",
    "    stride=256,\n",
    "    buffer_radius=0,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initial Training (First 10 epochs)\n",
    "\n",
    "First, we'll train a model for just 10 epochs and save checkpoints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initial training - just 10 epochs\n",
    "print(\"Starting initial training for 10 epochs...\")\n",
    "geoai.train_segmentation_model(\n",
    "    images_dir=f\"{out_folder}/images\",\n",
    "    labels_dir=f\"{out_folder}/labels\",\n",
    "    output_dir=f\"{out_folder}/initial_training\",\n",
    "    architecture=\"unet\",\n",
    "    encoder_name=\"resnet34\",\n",
    "    encoder_weights=\"imagenet\",\n",
    "    num_channels=3,\n",
    "    num_classes=2,\n",
    "    batch_size=4,  # Smaller batch size for faster testing\n",
    "    num_epochs=10,  # Just 10 epochs for testing\n",
    "    learning_rate=0.001,\n",
    "    val_split=0.2,\n",
    "    save_best_only=False,  # Save checkpoints every 10 epochs\n",
    "    verbose=True,\n",
    "    plot_curves=True,\n",
    ")\n",
    "print(\"Initial training completed!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Resume Training from Checkpoint\n",
    "\n",
    "Now we'll resume training from the checkpoint, continuing for another 10 epochs (total 20 epochs)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check if checkpoint exists\n",
    "checkpoint_path = f\"{out_folder}/initial_training/checkpoint_epoch_10.pth\"\n",
    "if os.path.exists(checkpoint_path):\n",
    "    print(f\"Found checkpoint: {checkpoint_path}\")\n",
    "else:\n",
    "    print(f\"Checkpoint not found: {checkpoint_path}\")\n",
    "    print(\"Available files:\")\n",
    "    for f in os.listdir(f\"{out_folder}/initial_training\"):\n",
    "        print(f\"  {f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Resume training from checkpoint\n",
    "print(\"Resuming training from checkpoint...\")\n",
    "geoai.train_segmentation_model(\n",
    "    images_dir=f\"{out_folder}/images\",\n",
    "    labels_dir=f\"{out_folder}/labels\",\n",
    "    output_dir=f\"{out_folder}/resumed_training\",\n",
    "    architecture=\"unet\",\n",
    "    encoder_name=\"resnet34\",\n",
    "    encoder_weights=\"imagenet\",\n",
    "    num_channels=3,\n",
    "    num_classes=2,\n",
    "    batch_size=4,\n",
    "    num_epochs=20,  # Total epochs (will resume from epoch 10)\n",
    "    learning_rate=0.001,\n",
    "    val_split=0.2,\n",
    "    save_best_only=False,\n",
    "    verbose=True,\n",
    "    plot_curves=True,\n",
    "    checkpoint_path=checkpoint_path,\n",
    "    resume_training=True,  # Resume training state\n",
    ")\n",
    "print(\"Resumed training completed!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Model Weights Only (No Training State Resume)\n",
    "\n",
    "Finally, we'll test loading only the model weights without resuming the training state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load weights only (no training state)\n",
    "print(\"Loading model weights only (not resuming training state)...\")\n",
    "geoai.train_segmentation_model(\n",
    "    images_dir=f\"{out_folder}/images\",\n",
    "    labels_dir=f\"{out_folder}/labels\",\n",
    "    output_dir=f\"{out_folder}/weights_only_training\",\n",
    "    architecture=\"unet\",\n",
    "    encoder_name=\"resnet34\",\n",
    "    encoder_weights=\"imagenet\",\n",
    "    num_channels=3,\n",
    "    num_classes=2,\n",
    "    batch_size=4,\n",
    "    num_epochs=15,  # Train for 15 epochs starting from epoch 0\n",
    "    learning_rate=0.001,\n",
    "    val_split=0.2,\n",
    "    save_best_only=False,\n",
    "    verbose=True,\n",
    "    plot_curves=True,\n",
    "    checkpoint_path=checkpoint_path,\n",
    "    resume_training=False,  # Don't resume training state\n",
    ")\n",
    "print(\"Weights-only training completed!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Best Model for Inference\n",
    "\n",
    "Test loading the best model for inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test inference with resumed model\n",
    "masks_path = f\"{out_folder}/test_prediction.tif\"\n",
    "model_path = f\"{out_folder}/resumed_training/best_model.pth\"\n",
    "\n",
    "if os.path.exists(model_path):\n",
    "    print(f\"Running inference with model: {model_path}\")\n",
    "    geoai.semantic_segmentation(\n",
    "        input_path=test_raster_path,\n",
    "        output_path=masks_path,\n",
    "        model_path=model_path,\n",
    "        architecture=\"unet\",\n",
    "        encoder_name=\"resnet34\",\n",
    "        num_channels=3,\n",
    "        num_classes=2,\n",
    "        window_size=512,\n",
    "        overlap=256,\n",
    "        batch_size=4,\n",
    "    )\n",
    "    print(f\"Inference completed. Results saved to: {masks_path}\")\n",
    "else:\n",
    "    print(f\"Model not found: {model_path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compare Training Histories\n",
    "\n",
    "Let's compare the training histories to verify that resuming worked correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Load training histories\n",
    "initial_history_path = f\"{out_folder}/initial_training/training_history.pth\"\n",
    "resumed_history_path = f\"{out_folder}/resumed_training/training_history.pth\"\n",
    "weights_only_history_path = f\"{out_folder}/weights_only_training/training_history.pth\"\n",
    "\n",
    "histories = {}\n",
    "for name, path in [\n",
    "    (\"Initial (10 epochs)\", initial_history_path),\n",
    "    (\"Resumed (10→20 epochs)\", resumed_history_path),\n",
    "    (\"Weights Only (15 epochs)\", weights_only_history_path),\n",
    "]:\n",
    "    if os.path.exists(path):\n",
    "        histories[name] = torch.load(path)\n",
    "        print(f\"Loaded {name}: {len(histories[name]['train_losses'])} epochs\")\n",
    "    else:\n",
    "        print(f\"History not found: {path}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot comparison - showing the continuation clearly\n",
    "if histories:\n",
    "    fig, axes = plt.subplots(1, 3, figsize=(18, 5))\n",
    "\n",
    "    # Get the individual histories\n",
    "    initial_history = histories[\"Initial (10 epochs)\"]\n",
    "    resumed_history = histories[\"Resumed (10→20 epochs)\"]\n",
    "    weights_only_history = histories[\"Weights Only (15 epochs)\"]\n",
    "\n",
    "    # Training Loss - show continuation clearly\n",
    "    initial_epochs = range(1, 11)\n",
    "    axes[0].plot(\n",
    "        initial_epochs,\n",
    "        initial_history[\"train_losses\"],\n",
    "        label=\"Initial (epochs 1-10)\",\n",
    "        marker=\"o\",\n",
    "        markersize=5,\n",
    "        color=\"blue\",\n",
    "        linewidth=3,\n",
    "        alpha=0.9,\n",
    "        zorder=3,\n",
    "    )\n",
    "\n",
    "    # Plot resumed training continuation (epochs 11-20 only)\n",
    "    resumed_epochs = range(11, 21)\n",
    "    resumed_continuation = resumed_history[\"train_losses\"][10:]  # epochs 11-20\n",
    "    axes[0].plot(\n",
    "        resumed_epochs,\n",
    "        resumed_continuation,\n",
    "        label=\"Resumed (epochs 11-20)\",\n",
    "        marker=\"s\",\n",
    "        markersize=5,\n",
    "        color=\"orange\",\n",
    "        linewidth=3,\n",
    "        alpha=0.9,\n",
    "        zorder=2,\n",
    "    )\n",
    "\n",
    "    # Plot weights only\n",
    "    weights_epochs = range(1, 16)\n",
    "    axes[0].plot(\n",
    "        weights_epochs,\n",
    "        weights_only_history[\"train_losses\"],\n",
    "        label=\"Weights Only (epochs 1-15)\",\n",
    "        marker=\"^\",\n",
    "        markersize=4,\n",
    "        color=\"green\",\n",
    "        linewidth=2,\n",
    "        alpha=0.7,\n",
    "        zorder=1,\n",
    "    )\n",
    "\n",
    "    # Add continuation line\n",
    "    axes[0].plot(\n",
    "        [10, 11],\n",
    "        [initial_history[\"train_losses\"][-1], resumed_continuation[0]],\n",
    "        color=\"red\",\n",
    "        linewidth=2,\n",
    "        linestyle=\"--\",\n",
    "        alpha=0.7,\n",
    "        label=\"Continuation\",\n",
    "    )\n",
    "\n",
    "    axes[0].set_title(\"Training Loss\", fontsize=14)\n",
    "    axes[0].set_xlabel(\"Epoch\")\n",
    "    axes[0].set_ylabel(\"Loss\")\n",
    "    axes[0].legend()\n",
    "    axes[0].grid(True, alpha=0.3)\n",
    "\n",
    "    # Validation IoU\n",
    "    axes[1].plot(\n",
    "        initial_epochs,\n",
    "        initial_history[\"val_ious\"],\n",
    "        label=\"Initial (epochs 1-10)\",\n",
    "        marker=\"o\",\n",
    "        markersize=5,\n",
    "        color=\"blue\",\n",
    "        linewidth=3,\n",
    "        alpha=0.9,\n",
    "        zorder=3,\n",
    "    )\n",
    "\n",
    "    resumed_iou_continuation = resumed_history[\"val_ious\"][10:]\n",
    "    axes[1].plot(\n",
    "        resumed_epochs,\n",
    "        resumed_iou_continuation,\n",
    "        label=\"Resumed (epochs 11-20)\",\n",
    "        marker=\"s\",\n",
    "        markersize=5,\n",
    "        color=\"orange\",\n",
    "        linewidth=3,\n",
    "        alpha=0.9,\n",
    "        zorder=2,\n",
    "    )\n",
    "\n",
    "    axes[1].plot(\n",
    "        weights_epochs,\n",
    "        weights_only_history[\"val_ious\"],\n",
    "        label=\"Weights Only (epochs 1-15)\",\n",
    "        marker=\"^\",\n",
    "        markersize=4,\n",
    "        color=\"green\",\n",
    "        linewidth=2,\n",
    "        alpha=0.7,\n",
    "        zorder=1,\n",
    "    )\n",
    "\n",
    "    axes[1].plot(\n",
    "        [10, 11],\n",
    "        [initial_history[\"val_ious\"][-1], resumed_iou_continuation[0]],\n",
    "        color=\"red\",\n",
    "        linewidth=2,\n",
    "        linestyle=\"--\",\n",
    "        alpha=0.7,\n",
    "        label=\"Continuation\",\n",
    "    )\n",
    "\n",
    "    axes[1].set_title(\"Validation IoU\", fontsize=14)\n",
    "    axes[1].set_xlabel(\"Epoch\")\n",
    "    axes[1].set_ylabel(\"IoU\")\n",
    "    axes[1].legend()\n",
    "    axes[1].grid(True, alpha=0.3)\n",
    "\n",
    "    # Validation Dice\n",
    "    axes[2].plot(\n",
    "        initial_epochs,\n",
    "        initial_history[\"val_dices\"],\n",
    "        label=\"Initial (epochs 1-10)\",\n",
    "        marker=\"o\",\n",
    "        markersize=5,\n",
    "        color=\"blue\",\n",
    "        linewidth=3,\n",
    "        alpha=0.9,\n",
    "        zorder=3,\n",
    "    )\n",
    "\n",
    "    resumed_dice_continuation = resumed_history[\"val_dices\"][10:]\n",
    "    axes[2].plot(\n",
    "        resumed_epochs,\n",
    "        resumed_dice_continuation,\n",
    "        label=\"Resumed (epochs 11-20)\",\n",
    "        marker=\"s\",\n",
    "        markersize=5,\n",
    "        color=\"orange\",\n",
    "        linewidth=3,\n",
    "        alpha=0.9,\n",
    "        zorder=2,\n",
    "    )\n",
    "\n",
    "    axes[2].plot(\n",
    "        weights_epochs,\n",
    "        weights_only_history[\"val_dices\"],\n",
    "        label=\"Weights Only (epochs 1-15)\",\n",
    "        marker=\"^\",\n",
    "        markersize=4,\n",
    "        color=\"green\",\n",
    "        linewidth=2,\n",
    "        alpha=0.7,\n",
    "        zorder=1,\n",
    "    )\n",
    "\n",
    "    axes[2].plot(\n",
    "        [10, 11],\n",
    "        [initial_history[\"val_dices\"][-1], resumed_dice_continuation[0]],\n",
    "        color=\"red\",\n",
    "        linewidth=2,\n",
    "        linestyle=\"--\",\n",
    "        alpha=0.7,\n",
    "        label=\"Continuation\",\n",
    "    )\n",
    "\n",
    "    axes[2].set_title(\"Validation Dice\", fontsize=14)\n",
    "    axes[2].set_xlabel(\"Epoch\")\n",
    "    axes[2].set_ylabel(\"Dice\")\n",
    "    axes[2].legend()\n",
    "    axes[2].grid(True, alpha=0.3)\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(\n",
    "        f\"{out_folder}/clear_training_comparison.png\", dpi=150, bbox_inches=\"tight\"\n",
    "    )\n",
    "    plt.show()\n",
    "\n",
    "    # Print explanation\n",
    "    print(\"\\n\" + \"=\" * 60)\n",
    "    print(\"CHECKPOINT RESUME VISUALIZATION EXPLANATION\")\n",
    "    print(\"=\" * 60)\n",
    "    print(\"• Blue line: Initial training (epochs 1-10)\")\n",
    "    print(\"• Orange line: Resumed training continuation (epochs 11-20)\")\n",
    "    print(\"• Green line: Weights-only training (fresh epochs 1-15)\")\n",
    "    print(\"• Red dashed line: Shows seamless continuation\")\n",
    "    print()\n",
    "    print(\"NOTE: In the original overlapping plot, the blue line was\")\n",
    "    print(\"completely covered by the orange line because resumed training\")\n",
    "    print(\"includes the exact same first 10 epochs plus 10 additional epochs.\")\n",
    "    print(\"=\" * 60)\n",
    "\n",
    "    # Print summary statistics\n",
    "    print(\"\\nTraining Summary:\")\n",
    "    for name, history in histories.items():\n",
    "        final_iou = history[\"val_ious\"][-1]\n",
    "        final_dice = history[\"val_dices\"][-1]\n",
    "        final_loss = history[\"train_losses\"][-1]\n",
    "        epochs_trained = len(history[\"train_losses\"])\n",
    "        print(f\"{name}: {epochs_trained} epochs\")\n",
    "        print(f\"  Final IoU: {final_iou:.4f}\")\n",
    "        print(f\"  Final Dice: {final_dice:.4f}\")\n",
    "        print(f\"  Final Loss: {final_loss:.4f}\")\n",
    "        print()\n",
    "\n",
    "    print(\"Training comparison plot saved to clear_training_comparison.png\")\n",
    "else:\n",
    "    print(\"No training histories found for comparison\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Figure_1](https://github.com/user-attachments/assets/1560e0a7-1e16-449c-97de-8362bc755f46)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "This notebook demonstrates the new checkpoint functionality:\n",
    "\n",
    "1. **Initial Training**: Trained for 10 epochs and saved checkpoints\n",
    "2. **Resume Training**: Successfully resumed from epoch 10 and continued to epoch 20\n",
    "3. **Weights Only**: Loaded model weights but started training from epoch 0\n",
    "4. **Inference**: Used the resumed model for inference\n",
    "\n",
    "The key features tested:\n",
    "- `checkpoint_path`: Path to the checkpoint file\n",
    "- `resume_training=True`: Resume training state (epoch, optimizer, scheduler, metrics)\n",
    "- `resume_training=False`: Load only model weights\n",
    "\n",
    "This functionality is useful for:\n",
    "- Long training jobs that might be interrupted\n",
    "- Experimenting with different training parameters after initial training\n",
    "- Transfer learning from partially trained models"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "geo",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
